import numpy as np
import cv2

def project_perspective(pano, center, width, height, f):
    '''
    全景图生成中心投影 TODO lookat有bug 待修复
    :param pano:
    :param center:
    :param width:
    :param height:
    :param f:
    :return:
    '''
    h, w, _ = pano.shape
    assert w == 2 * h
    proj = np.zeros((height, width, 3), np.uint8)
    cx, cy = width / 2, height / 2
    x = np.arange(width)
    y = np.arange(height)
    theta = np.arctan2(x - cx, f)
    for i in range(width):
        u = np.rint(theta[i] / np.pi * h + center[1]).astype(int)
        tan_phi = (y - cy) / np.sqrt((x[i] - cx) ** 2 + f ** 2)
        v = np.rint(tan_phi / 2 * h + center[0]).astype(int)
        # phi = np.arctan2((y - cy), np.sqrt((x[i] - cx) ** 2 + f ** 2))
        # v = np.rint(phi /np.pi * h  + center[0]).astype(int)
        proj[:, x[i]] = pano[v, u]
        pano[v, u] = [0, 255, 0]
    return proj


def project_pano(pcd, pano):
    '''
    点云投影成全景图
    :param pcd: 在相机坐标系下的点云
    :param pano: 对应的全景图片
    :return: 点云投影的图片（二值）
    '''
    h, w, _ = pano.shape
    assert w == 2 * h
    xyz = np.asarray(pcd.points)
    rtp = np.ndarray(xyz.shape)
    xy = xyz[:, 0] ** 2 + xyz[:, 1] ** 2
    rtp[:, 0] = np.sqrt(xy + xyz[:, 2] ** 2) # rho
    rtp[:, 1] = w - np.mod(2 * np.pi + np.arctan2(xyz[:, 1], xyz[:, 0]), 2 * np.pi) * h / np.pi # theta
    rtp[:, 2] = -xyz[:, 2]/np.sqrt(xy) * (h / 2) + (h / 2) # phi
    prj_pano = np.zeros((int(pano.shape[0] / 4), int(pano.shape[1] / 4)))
    for i, p in enumerate(rtp):
        cv2.circle(pano, (int(p[1]), int(p[2])), 4, (0, 0, 255), -1)  # 在原图片上画点
        cv2.circle(prj_pano, (int(p[1] / 4), int(p[2] / 4)), 4, 255, -1)  # 在空白上图片画点
    # cv2.imwrite("./data/proj.png",prj_pano)
    return prj_pano


if __name__ == "__main__":
    pano = cv2.imread(r"C:\Users\liyif\Documents\ladybug_20210320_104703_174168_Panoramic_8192x4096.jpg")
    cv2.namedWindow("a", cv2.WINDOW_KEEPRATIO + cv2.WINDOW_NORMAL)

    proj = project_perspective(pano, (2048, 1024 * 2), 1024, 768, 400)

    cv2.imshow("a", pano)
    cv2.imshow("proj", proj)
    cv2.waitKey(0)
